import numpy as np
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt
import pandas as pd
from python_ai.common.xcommon import sep
from python_ai.ML_2.decomposition.follow_teacher.x_pca_impl_again import x_cross_check_us, x_pca_impl_again, x_pca_impl_svd_again


def x_pca_by_svd(X, k):
    """
    https://numpy.org/doc/stable/reference/generated/numpy.cov.html
    https://numpy.org/doc/stable/reference/generated/numpy.linalg.svd.html?highlight=svd#numpy.linalg.svd

    :param X:
    :param k:
    :return:
    """
    m, n = X.shape

    # 去中心化
    mu = X.mean(axis=0)
    X -= mu

    # 协方差矩阵
    cov_mat = np.cov(X.T)  # ATTENTION X.T

    # 特征值lmds、特征向量us
    us, lmds, _ = np.linalg.svd(cov_mat)

    # 特征值、特征向量选取
    lmds = lmds[:k]
    u = us[:, :k]  # ATTENTION us[:, :k]

    # decomposition
    x_new = X.dot(u)
    return x_new, lmds, u.T  # ATTENTION us.T


def x_pca(X, k):
    """
    PCA

    :param X: Feature matrix.
    :param k: Target dimension number.
    :return:
    """
    m, n = X.shape

    # 去中心化
    mu = X.mean(axis=0)
    X -= mu

    # 协方差矩阵
    cov_mat = np.dot(X.T, X) / m

    # 特征值lmds、特征向量us
    lmds, us = np.linalg.eig(cov_mat)

    # 特征值、特征向量排序、选取
    eig_pair = []
    for i in range(n):
        pair = (lmds[i], us[i])
        eig_pair.append(pair)
    eig_pair.sort(key=lambda x: x[0], reverse=True)
    eig_pair = np.array(eig_pair)
    u = eig_pair[0][1]
    for i in range(1, k):
        u = np.c_[u, eig_pair[i][1]]

    # decomposition
    x_new = X.dot(u)
    return x_new, eig_pair[:k, 0], u.T


if '__main__' == __name__:

    X = np.array([
        [1, 2, 3],
        [4, 5, 6],
        [7, 7, 9],
        [2, 4, 6],
        [3, 6, 9]
    ], dtype=np.float64)

    fig = plt.figure(figsize=[16, 9])
    spr = 2
    spc = 4
    spn = 0

    # plot 1
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.scatter3D(X[:, 0], X[:, 1], X[:, 2])
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    # plot 2
    title = 'pca => 2 dim by self impl'
    sep(title)
    X2, lmds, us = x_pca(X, 2)
    # X2, lmds, us = x_pca_impl_again(X, 2)
    print(X2)
    print(lmds)
    print(us)
    x_cross_check_us(us)
    cmap = plt.cm.get_cmap('rainbow', len(us))
    for i, ui in enumerate(us):
        ui *= 10
        ax.plot3D([0, ui[0]], [0, ui[1]], [0, ui[2]], c=cmap(i))
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.scatter(X2[:, 0], X2[:, 1])

    # plot 3
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.scatter3D(X[:, 0], X[:, 1], X[:, 2])
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    # plot 4
    title = 'pca => 2 dim by sklearn PCA'
    sep(title)
    from sklearn.decomposition import PCA
    dc = PCA(n_components=2)
    X2 = dc.fit_transform(X)
    lmds = dc.explained_variance_
    us = dc.components_
    print(X2)
    print(lmds)
    print(us)
    x_cross_check_us(us)
    for i, ui in enumerate(us):
        ui *= 10
        ax.plot3D([0, ui[0]], [0, ui[1]], [0, ui[2]], c=cmap(i))
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.scatter(X2[:, 0], X2[:, 1])

    # plot 5
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.scatter3D(X[:, 0], X[:, 1], X[:, 2])
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    # plot 6
    title = 'pca => 2 dim by TruncatedSVD'
    sep(title)
    from sklearn.decomposition import TruncatedSVD
    dc = TruncatedSVD(n_components=2)
    X_m = X - X.mean(axis=0)  # ATTENTION
    X2 = dc.fit_transform(X_m)
    lmds = dc.explained_variance_
    us = dc.components_
    print(X2)
    print(lmds)
    print(us)
    x_cross_check_us(us)
    for i, ui in enumerate(us):
        ui *= 10
        ax.plot3D([0, ui[0]], [0, ui[1]], [0, ui[2]], c=cmap(i))
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.scatter(X2[:, 0], X2[:, 1])

    # plot 7
    spn += 1
    ax = plt.subplot(spr, spc, spn, projection='3d')
    ax.scatter3D(X[:, 0], X[:, 1], X[:, 2])
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    # plot 8
    title = 'pca => 2 dim by self impl based on svd'
    sep(title)
    X2, lmds, us = x_pca_by_svd(X, 2)
    # X2, lmds, us = x_pca_impl_svd_again(X, 2)
    print(X2)
    print(lmds)
    print(us)
    x_cross_check_us(us)
    for i, ui in enumerate(us):
        ui *= 10
        ax.plot3D([0, ui[0]], [0, ui[1]], [0, ui[2]], c=cmap(i))
    spn += 1
    plt.subplot(spr, spc, spn)
    plt.title(title)
    plt.scatter(X2[:, 0], X2[:, 1])

    plt.show()
